import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
import matplotlib.pyplot as plt


def main():
    """读入数据"""
    df_wine = pd.read_csv('./wine/wine.data', header=None)  # 本地加载

    """取其中类别标签为一类和二类的白酒数据生成新的白酒数据"""
    df_wine = df_wine.drop(index=df_wine[(df_wine[0] == 3)].index.tolist())

    # 把分类属性与常规属性分开
    xa, y = df_wine.iloc[:, 1:].values, df_wine.iloc[:, 0].values
    # 有num个种类
    # 降到k维
    num = 2
    k = 2

    # 所有元素的平均值
    ua = np.array([np.mean(xa, axis=0)])
    n = xa.shape[1]
    # 存储x值
    x = [[]for i in range(num)]
    u = [[]for i in range(num)]
    sw = np.zeros([n, n])
    sb = np.zeros([n, n])
    for i in range(y.shape[0]):
        x[y[i]-1].append(xa[i])
    for i in range(num):
        x[i] = np.array(x[i])
        u[i] = np.array([np.mean(x[i], axis=0)])
    # 去中心化,计算Sw
    for i in range(num):
        x[i] = x[i]-u[i]
        sw = sw+np.dot(x[i].T, x[i])
        sb = sb+x[i].shape[0]*np.dot((u[i]-ua).T, (u[i]-ua))
    # 求Sw_1Sb的特征值与特征向量
    eig_value, eig_vec = np.linalg.eig(np.dot(np.linalg.inv(sw), sb))
    print(eig_vec)
    # 对特征值进行排序，返回标号
    index = np.argsort(eig_value)
    # 获取取最后k个特征向量
    eig_vec_k = eig_vec[:, index[:-k-1:-1]]
    colors = ["r", "g", "b"]

    # 绘图
    for i in range(xa.shape[0]):
        plt.scatter(np.dot(xa[i], eig_vec_k[:, 0]), np.dot(
            xa[i], eig_vec_k[:, 1]), color=colors[y[i]], label=y[i])
    plt.title("my LDA")
    plt.show()


if __name__ == '__main__':
    main()
